ATOM Documentation

┌─────────────────────────────────────────────────────────────┐
│ SaaS Architecture                                           │
├─────────────────────────────────────────────────────────────┤
│ LLMService (730 lines)                                      │
│   ├── Unified API for generation, completion, embeddings    │
│   ├── Continuous learning personalization                   │
│   ├── Token estimation & cost tracking                      │
│   └── Wraps BYOKHandler                                     │
├─────────────────────────────────────────────────────────────┤
│ BYOKHandler (2,064 lines)                                   │
│   ├── Cognitive tier classification                         │
│   ├── Cache-aware router                                    │
│   ├── BPC provider ranking                                  │
│   ├── Circuit breaker & retry                               │
│   └── Provider health monitoring                            │
├─────────────────────────────────────────────────────────────┤
│ BYOKManager (1,437 lines)                                   │
│   ├── Multi-tenant API key storage (encrypted)              │
│   ├── Provider configuration                                │
│   └── Usage tracking per tenant                             │
└─────────────────────────────────────────────────────────────┘

┌─────────────────────────────────────────────────────────────┐
│ Open-Source Architecture                                    │
├─────────────────────────────────────────────────────────────┤
│ BYOKHandler (1,839 lines)                                   │
│   ├── Cognitive tier classification                         │
│   ├── Cache-aware router                                    │
│   ├── BPC provider ranking                                  │
│   ├── Circuit breaker & retry                               │
│   └── Provider health monitoring                            │
├─────────────────────────────────────────────────────────────┤
│ BYOKManager (1,297 lines)                                   │
│   ├── Single-tenant API key storage (encrypted)             │
│   ├── Provider configuration                                │
│   └── Usage tracking                                        │
├─────────────────────────────────────────────────────────────┤
│ CognitiveTierService (526 lines)                            │
│   ├── Orchestration layer for tier routing                  │
│   ├── Workspace preference management                       │
│   └── Budget constraint checking                            │
└─────────────────────────────────────────────────────────────┘

class LLMService:
    # Text Generation
    async def generate(...) -> str
    async def generate_completion(...) -> Dict[str, Any]
    async def generate_structured_response(...) -> Any
    async def stream_completion(...) -> AsyncGenerator[str, None]
    
    # Embeddings
    async def generate_embedding(...) -> List[float]
    async def generate_embeddings_batch(...) -> List[List[float]]
    
    # Multimodal
    async def transcribe_audio(...) -> Dict[str, Any]
    async def generate_speech(...) -> bytes
    
    # Cognitive Tier Routing
    async def generate_with_tier(...) -> Dict[str, Any]
    def get_optimal_provider(...) -> tuple[str, str]
    def get_ranked_providers(...) -> List[tuple[str, str]]
    
    # Utilities
    def estimate_tokens(...) -> int
    def estimate_cost(...) -> float

# SaaS pattern - via LLMService wrapper
llm_service = LLMService(db=session, workspace_id="ws-123", tenant_id="tenant-456")
response = await llm_service.generate(
    prompt="Analyze this data...",
    model="auto",  # Auto-routed by cognitive tier
    temperature=0.7,
    agent_id="agent-789",  # Enables personalization
    tenant_id="tenant-456"
)

# Open-Source pattern - direct BYOKHandler usage
handler = BYOKHandler(workspace_id="ws-123", db_session=session)
response = await handler.generate_response(
    prompt="Analyze this data...",
    model_type="auto",
    temperature=0.7
)

if agent_id and self.continuous_learning:
    params = self.continuous_learning.get_personalized_parameters(
        tenant_id=target_ws,
        agent_id=agent_id,
        user_id=user_id
    )
    if "temperature" in params:
        temperature = params["temperature"]

llm_usage_tracker.record(
    workspace_id=target_ws,
    provider=provider,
    model=model,
    input_tokens=input_tokens,
    output_tokens=output_tokens,
    cost_usd=cost,
    user_id=user_id,
    agent_id=agent_id,
    is_managed_service=kwargs.get("is_managed_service", False),
    chain_id=kwargs.get("chain_id")
)

[
    "openai",        # GPT-5.3, GPT-4o
    "anthropic",     # Claude 4.6 Opus, Claude 3.5 Sonnet
    "moonshot",      # Kimi k1.5 Thinking
    "google",        # Gemini 1.5 Pro
    "google_flash",  # Gemini 1.5 Flash
    "lux",           # LUX Computer Use
    "deepseek",      # DeepSeek-V3, DeepSeek-R1
    "glm",           # GLM-4, GLM-4.6, GLM-5
    "minimax",       # MiniMax M2.7
    "qwen",          # Qwen-Max, Qwen-Plus
    "deepinfra"      # Open-source models
]

[
    "deepseek",      # DeepSeek-V3 (primary)
    "openai",        # GPT-4o, GPT-3.5
    "anthropic",     # Claude 3.5 Sonnet
    "groq",          # Llama 3.3/3.1
    "google",        # Gemini 1.5 Pro
    "google_flash",  # Gemini 1.5 Flash
    "minimax",       # MiniMax M2.5
    "moonshot",      # Kimi
    "deepinfra"      # Open-source models
]

def _encrypt_key(self, api_key: str) -> str:
    fernet = Fernet(self.encryption_key)
    return fernet.encrypt(api_key.encode()).decode()

def _decrypt_key(self, encrypted_key: str) -> str:
    fernet = Fernet(self.encryption_key)
    return fernet.decrypt(encrypted_key.encode()).decode()

def get_ranked_providers(self, complexity, ...):
    for model_id, pricing in fetcher.pricing_cache.items():
        # 1. Filter by context window
        if context_window < min_context:
            continue
        
        # 2. Filter by quality score (CognitiveTier thresholds)
        if quality_score < min_quality:
            continue
        
        # 3. Filter by capabilities (tools, vision, etc.)
        if required_capability and required_capability not in capabilities:
            continue
        
        # 4. Calculate cache-aware effective cost
        effective_cost = cache_router.calculate_effective_cost(
            model=model_id,
            provider=active_provider,
            estimated_input_tokens=estimated_tokens,
            cache_hit_probability=0.5
        )
        
        # 5. Compute value score
        if prefer_cost:
            value_score = quality_score / (effective_cost + 1e-9)
        else:
            value_score = quality_score * (1.0 / (effective_cost + 1e-9))
        
        ranked_options.append((value_score, active_provider, model_id))
    
    # Sort by value score descending
    ranked_options.sort(reverse=True, key=lambda x: x[0])
    return [(provider, model) for _, provider, model in ranked_options]

class CognitiveTier(Enum):
    MICRO = "micro"       # Simple greetings, <50 tokens
    STANDARD = "standard" # Basic Q&A, 50-500 tokens
    VERSATILE = "versatile" # Analysis, 500-2000 tokens
    HEAVY = "heavy"       # Complex reasoning, 2000-5000 tokens
    COMPLEX = "complex"   # Expert tasks, 5000+ tokens

def classify(self, prompt: str, task_type: Optional[str] = None) -> CognitiveTier:
    # 1. Length-based scoring
    estimated_tokens = len(prompt) / 4
    if estimated_tokens >= 5000: score += 4
    elif estimated_tokens >= 2000: score += 3
    elif estimated_tokens >= 500: score += 2
    elif estimated_tokens >= 50: score += 1
    
    # 2. Keyword analysis
    patterns = {
        "simple": (r"\b(hello|hi|thanks|summarize|list)\b", -1),
        "moderate": (r"\b(analyze|compare|explain|describe)\b", 1),
        "technical": (r"\b(calculate|solve|equation|code|debug)\b", 2),
        "advanced": (r"\b(architecture|security|distributed|optimize)\b", 3)
    }
    
    # 3. Task type override
    if task_type == "code": score += 1
    if task_type == "chat": score -= 1
    
    # 4. Map to tier
    if score <= 0: return MICRO
    elif score == 1: return STANDARD
    elif score == 2: return VERSATILE
    elif score == 3: return HEAVY
    else: return COMPLEX

CACHE_CAPABILITIES = {
    "openai": {
        "supports_cache": True,
        "cached_cost_ratio": 0.10,  # 90% discount
        "min_tokens": 1024,
    },
    "anthropic": {
        "supports_cache": True,
        "cached_cost_ratio": 0.10,
        "min_tokens": 2048,  # Longer prompts required
    },
    "gemini": {
        "supports_cache": True,
        "cached_cost_ratio": 0.10,
        "min_tokens": 1024,
    },
    "deepseek": {
        "supports_cache": False,  # No caching
        "cached_cost_ratio": 1.0,
        "min_tokens": 0,
    },
    "minimax": {
        "supports_cache": False,
        "cached_cost_ratio": 1.0,
        "min_tokens": 0,
    },
}

def calculate_effective_cost(
    self,
    model: str,
    provider: str,
    estimated_input_tokens: int,
    cache_hit_probability: float = 0.5
) -> float:
    # Get list price
    input_cost = pricing.get("input_cost_per_token", 0)
    output_cost = pricing.get("output_cost_per_token", 0)
    
    # Check cache capability
    cache_info = self.get_provider_cache_capability(provider)
    if not cache_info["supports_cache"]:
        return (input_cost + output_cost) / 2  # Full price
    
    # Check minimum token threshold
    if estimated_input_tokens < cache_info["min_tokens"]:
        return (input_cost + output_cost) / 2  # Too short for caching
    
    # Calculate effective cost with cache hit probability
    cached_ratio = cache_info["cached_cost_ratio"]
    effective_input_cost = input_cost * (
        cache_hit_probability * cached_ratio +      # Cached portion
        (1 - cache_hit_probability) * 1.0           # Uncached portion
    )
    
    return (effective_input_cost + output_cost) / 2

GET  /byok/keys?tenant_id=...          # List tenant's provider keys
POST /byok/keys?tenant_id=...          # Add new API key
DELETE /byok/keys/{provider_id}?tenant_id=...  # Remove key

GET  /api/llm-registry/provider-health?providers=...  # Provider health status
GET  /api/llm-registry/models/by-quality?min_quality=80&capabilities=...  # Filter by quality
POST /api/llm-registry/sync-quality?source=lmsys&force_refresh=false  # Sync quality scores

GET    /api/v1/cognitive-tier/preferences/{workspace_id}  # Get tier preferences
POST   /api/v1/cognitive-tier/preferences/{workspace_id}  # Set preferences
PUT    /api/v1/cognitive-tier/preferences/{workspace_id}/budget  # Update budget
GET    /api/v1/cognitive-tier/estimate-cost?prompt=...&estimated_tokens=100  # Cost estimate

POST /api/v1/byok/add-key              # Add API key (secure POST body)
GET  /api/v1/byok/providers            # List available providers
GET  /api/v1/byok/usage/{provider_id}  # Get usage stats

llm_usage_tracker.record(
    workspace_id="ws-123",
    provider="deepseek",
    model="deepseek-chat",
    input_tokens=1500,
    output_tokens=500,
    cost_usd=0.00035,
    user_id="user-456",
    agent_id="agent-789",
    is_managed_service=True,
    chain_id="chain-abc"
)

# Cost-optimized ranking
value_score = quality_score / (effective_cost + 1e-9)

# Quality-optimized ranking
value_score = quality_score * (1.0 / (effective_cost + 1e-9))

class ProviderHealthService:
    # Tracks per-provider:
    # - Success rate (last 1000 requests)
    # - Error rate (last 1000 requests)
    # - Consecutive failures
    # - Average latency
    # - Rate limit status
    
    # Circuit breaker states:
    # - HEALTHY (success_rate >= 95%)
    # - DEGRADED (success_rate 80-95%)
    # - UNHEALTHY (success_rate < 80%)
    # - RATE_LIMITED (429 responses)

capabilities = ["chat", "code", "vision", "tools", "structured_output", "computer_use"]

# Get models with tool calling
get_models_by_quality_range(db, tenant_id, min_quality=80, capabilities=["tools"])

# Get vision-capable models
get_models_by_quality_range(db, tenant_id, min_quality=86, capabilities=["vision"])

from core.byok_endpoints import get_byok_manager

byok_manager = get_byok_manager()

# Store tenant-specific key
key_id = byok_manager.store_tenant_api_key(
    tenant_id="tenant-456",
    provider_id="deepseek",
    api_key="sk-...",
    key_name="production",
    db=db_session
)

# Retrieve tenant key
api_key = byok_manager.get_tenant_api_key(
    tenant_id="tenant-456",
    provider_id="deepseek",
    db=db_session
)

# Delete tenant key
byok_manager.delete_tenant_api_key(
    tenant_id="tenant-456",
    provider_id="deepseek",
    db=db_session
)

from api.cognitive_tier_routes import router

# Get workspace preferences
# GET /api/v1/cognitive-tier/preferences/ws-123
# Response:
{
    "workspace_id": "ws-123",
    "default_tier": "versatile",
    "min_tier": "standard",
    "max_tier": "heavy",
    "monthly_budget_cents": 5000,
    "per_request_budget_cents": 50
}

# Set preferences
# POST /api/v1/cognitive-tier/preferences/ws-123
{
    "default_tier": "versatile",
    "min_tier": "standard",
    "max_tier": "heavy",
    "monthly_budget_cents": 5000
}

# Update budget
# PUT /api/v1/cognitive-tier/preferences/ws-123/budget
{
    "monthly_budget_cents": 10000,
    "per_request_budget_cents": 100
}

# SaaS via LLMService
from core.llm_service import LLMService

llm = LLMService(db=db_session, workspace_id="ws-123")
response = await llm.generate_with_tier(
    prompt="Analyze this distributed system architecture...",
    system_instruction="You are a senior software architect.",
    task_type="analysis",
    agent_id="agent-789"
)
# Returns: {"response": "...", "tier_used": "heavy", "model": "claude-4-opus", "cost_cents": 2.5}

# Open-Source via BYOKHandler
from core.llm.byok_handler import BYOKHandler

handler = BYOKHandler(workspace_id="ws-123", db_session=db_session)
response = await handler.generate_with_cognitive_tier(
    prompt="Analyze this distributed system architecture...",
    system_instruction="You are a senior software architect.",
    task_type="analysis"
)
# Returns same structure

tests/
├── test_byok_logic.py              # BYOKManager unit tests
├── test_llm_service.py             # LLMService wrapper tests
├── test_cognitive_tier_routing.py  # Tier classification tests
└── api/security/test_byok_security.py  # Encryption & isolation tests

tests/
├── test_cognitive_tier_classification.py  # Tier classification + BYOK integration
├── test_llm_endpoints_integration.py      # Full endpoint integration tests
├── test_pdf_ocr_vision.py                 # Vision model tests
└── test_byok_cost_optimizer.py            # Cost optimization tests

# Per-provider rate limits
max_requests_per_minute: int = 60
rate_limit_window: int = 60  # seconds

# Tracked per tenant (SaaS) or globally (Open-Source)
rate_limit_remaining: int
rate_limit_reset: Optional[datetime]

backend-saas/
├── core/
│   ├── llm_service.py                    # LLMService wrapper
│   ├── byok_endpoints.py                 # BYOKManager (1,437 lines)
│   └── llm/
│       ├── byok_handler.py               # BYOKHandler (2,064 lines)
│       ├── cognitive_tier_service.py     # Orchestration layer
│       ├── cognitive_tier_system.py      # Tier classification
│       ├── cache_aware_router.py         # Cache optimization
│       ├── registry/
│       │   ├── provider_health.py        # Health monitoring
│       │   └── queries.py                # Model filtering
│       └── fallback/
│           ├── circuit_breaker.py        # Resilience pattern
│           └── retry_policy.py           # Retry logic
└── api/
    ├── byok_api_routes.py                # BYOK management
    └── routes/
        └── llm_registry_routes.py        # Registry endpoints

atom-upstream/backend/
├── core/
│   ├── byok_endpoints.py                 # BYOKManager (1,297 lines)
│   └── llm/
│       ├── byok_handler.py               # BYOKHandler (1,839 lines)
│       ├── cognitive_tier_service.py     # Orchestration layer
│       ├── cognitive_tier_system.py      # Tier classification
│       ├── cache_aware_router.py         # Cache optimization
│       └── escalation_manager.py         # Quality-based escalation
└── api/
    ├── cognitive_tier_routes.py          # Tier preference API
    └── routes/
        └── byok_routes.py                # BYOK management (if exists)